package main

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io/ioutil"
)

type TLSMsgBody struct {
	DataSec   []byte //  数据密文
	KeySec    []byte // 对称密钥密文
	signature []byte // 摘要的签名
}

const (
	SERVER int8 = iota
	Client
)

// 完整的加密传输流程
//
// 发送端
// 1 将传输的文件进行hash得到摘要（消息完整性）
// 1 用接受端的公钥对传输的文件进行加密得到密文（对称加密）
// 2 使用自己的私钥对摘要进行加密得到数字签名（身份认证）
// 3 将传输的文件和数字签名一起发送给接收端
//
// 接收端
// 1 使用自己的私钥解密密文得到传输文件的明文
// 2 使用发送端的公钥对数字签名解密，得到正确的摘要（解密成功也代表身份验证成功）
// 3 对明文进行摘要计算，得到实际收到的摘要，比较两份摘要，如果一致，说明消息未被篡改（消息完整性）
func main() {

	sender := Sender()

	Receiver(sender)

}

func Sender() *TLSMsgBody {
	var sendData = new(TLSMsgBody)
	// 明文内容
	data, err := ioutil.ReadFile("./verify/hello.txt")
	if err != nil {
		fmt.Println(err)
		return nil
	}
	// 明文加密
	// 对称加密密钥
	// 16,24,32位字符串的话，分别对应AES-128，AES-192，AES-256 加密方法
	var pwdKey = []byte("ABCDABCDABCDABCD")
	aes := &Aes{
		PwdKey: pwdKey,
	}
	sendData.DataSec, err = aes.AesEncrypt(data, pwdKey)
	if err != nil {
		fmt.Println(err)
		return nil
	}

	// 将加密后的文件保存，实际使用中传输该文件
	if err = ioutil.WriteFile("./verify/hello_sec.txt", sendData.DataSec, 0644); err != nil {
		panic(err)
	}

	// 非对称加密
	recvPublic := initPubKey(SERVER)
	var r = Rsa{}
	encrypt, _ := r.Encrypt(recvPublic, pwdKey)
	sendData.KeySec = encrypt
	fmt.Printf("对称密钥：%s\n 对称加密完成=> %v\n", pwdKey, encrypt)

	// 摘要
	var m = Md5{}
	encode := m.Encode(data)
	fmt.Printf("摘要：%v\n", encode) //4bc3555484fca9f2879fe1a68b139e2d

	//digString := "4bc3555484fca9f2879fe1a68b139e2d"
	//check := m.Check(data, digString)
	//if check {
	//	fmt.Println("pass")
	//} else {
	//	fmt.Println("fail:file has been changed")
	//}

	// 摘要签名
	privateKey := initPriKey(Client)
	sendData.signature, err = rsa.SignPSS(rand.Reader, privateKey, crypto.MD5, encode, nil)
	if err != nil {
		panic(err)
	}
	fmt.Printf("摘要：%v 数字签名完成=> %v\n", encode, sendData.signature)
	return sendData
}

func Receiver(t *TLSMsgBody) {
	// 密文
	data, err := ioutil.ReadFile("./verify/hello_sec.txt")
	if err != nil {
		fmt.Println(err)
		return
	}

	KeySec := t.KeySec

	// 获取明文
	priKey := initPriKey(SERVER)
	var r = Rsa{}
	decrypt, err := r.Decrypt(priKey, KeySec)
	if err != nil {
		fmt.Println(err)
		return
	}

	aes := &Aes{
		PwdKey: decrypt,
	}
	// aesDecrypt对应的就是传输文件明文
	aesDecrypt, err := aes.AesDecrypt(data, decrypt)
	if err != nil {
		fmt.Println(err)
		return
	}
	// 解密后的文件
	if err = ioutil.WriteFile("./verify/hello_out.txt", aesDecrypt, 0644); err != nil {
		panic(err)
	}

	// 获取摘要签名
	signatureOld := t.signature
	pubKey := initPubKey(Client)

	// 计算接收文件的摘要
	var m = Md5{}
	encode := m.Encode(aesDecrypt)
	fmt.Printf("计算的摘要：%v\n", encode) //4bc3555484fca9f2879fe1a68b139e2d

	// m.check 比对摘要是否一致，可忽略
	check := m.Check(aesDecrypt, string(encode))
	fmt.Println("摘要比对结果，", check)

	// rsa.VerifyPSS验证签名是否一致
	err = rsa.VerifyPSS(pubKey, crypto.MD5, encode, signatureOld, nil)
	if err != nil {
		fmt.Println("签名比对不通过，", err)
		return
	}
	fmt.Println("文件完整")
}

func initPubKey(certType int8) *rsa.PublicKey {
	var (
		pubCert []byte
		err     error
	)
	switch certType {
	case Client:
		pubCert, err = ioutil.ReadFile("./client/c101.cert.pem")
		break
	case SERVER:
		pubCert, err = ioutil.ReadFile("./server/cadmus.cert.pem")
	}

	if err != nil {
		fmt.Println(err)
		return nil
	}
	publicKey, err := getPublicKeyFromCert(pubCert)
	if err != nil {
		fmt.Println(err)
		return nil
	}
	return publicKey
}

func initPriKey(certType int8) *rsa.PrivateKey {
	var (
		priPem []byte
		err    error
	)
	switch certType {
	case Client:
		priPem, err = ioutil.ReadFile("./client/c101.key.pem")
		break
	case SERVER:
		priPem, err = ioutil.ReadFile("./server/cadmus.key.pem")
	}
	if err != nil {
		fmt.Println(err)
		return nil
	}
	privateKey, err := getPrivateKey(priPem)
	if err != nil {
		fmt.Println(err)
		return nil
	}
	return privateKey
}

func getPublicKeyFromCert(cert []byte) (publicKey *rsa.PublicKey, err error) {
	block, _ := pem.Decode(cert)

	certificate, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		fmt.Println(err)
		return
	}

	publicKey = certificate.PublicKey.(*rsa.PublicKey)
	return
}

func getPrivateKey(cert []byte) (privateKey *rsa.PrivateKey, err error) {
	block, _ := pem.Decode(cert)

	privateKey, err = x509.ParsePKCS1PrivateKey(block.Bytes)
	return
}
